use anyhow::Context;
use aya::programs::{Xdp, XdpFlags};
use clap::Parser;
use compile::compile_dir;
use linker::linker;
use log::{debug, warn};
use network_scanner_detector_compiler::record::{RecordHandler, RecordItem};
use rpc::{config::parse_config, connect};
use std::{error::Error, sync::Arc};
use tokio::signal;

mod compile;
mod linker;
mod rpc;

#[derive(Debug, Parser)]
struct Opt {
    #[clap(short, long, default_value = "enp49s0")]
    interface: String,
    #[clap(short, long, default_value = "protos")]
    proto_dir: String,
    #[clap(short, long, default_value = "output")]
    output_dir: String,
    #[clap(short, long, default_value = "config.toml")]
    config: String,
}

struct RecordHandlerImpl {
    index: u16,
}

impl RecordHandlerImpl {
    fn new() -> Self {
        Self { index: 0 }
    }
}

impl RecordHandler for RecordHandlerImpl {
    fn emit_item(&mut self, item: &RecordItem) -> u16 {
        println!("ADD RECORD ITEM: {:?}", item);
        let index = self.index;
        self.index += 1;
        index
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let opt = Opt::parse();
    let mut record_handler = RecordHandlerImpl::new();
    let output_files = compile_dir(
        opt.proto_dir.as_str(),
        opt.output_dir.as_str(),
        &mut record_handler,
    )
    .await?;
    let object_path = format!("{}/main.o", opt.output_dir);
    let config = parse_config(&opt.config);
    let client = Arc::new(connect(config.rpc).await);
    linker(&output_files, &object_path)?;
    let rlim = libc::rlimit {
        rlim_cur: libc::RLIM_INFINITY,
        rlim_max: libc::RLIM_INFINITY,
    };
    let ret = unsafe { libc::setrlimit(libc::RLIMIT_MEMLOCK, &rlim) };
    if ret != 0 {
        debug!("remove limit on locked memory failed, ret is: {}", ret);
    }

    let mut ebpf = aya::Ebpf::load_file(object_path)?;
    println!("ebpf object loaded");
    if let Err(e) = aya_log::EbpfLogger::init(&mut ebpf) {
        warn!("failed to initialize eBPF logger: {}", e);
    }
    println!("ebpf logger initialized");
    let prog: &mut Xdp = ebpf.program_mut("runtime").unwrap().try_into()?;
    prog.load()?;
    println!("ebpf program loaded");
    prog.attach(&opt.interface, XdpFlags::default())
        .context("failed to attach the XDP program with default flags - try changing XdpFlags::default() to XdpFlags::SKB_MODE")?;

    let ctrl_c = signal::ctrl_c();
    println!("Waiting for Ctrl-C...");
    ctrl_c.await?;
    println!("Exiting...");

    Ok(())
}
